package org.example.provider.dubbo;

import brave.Tracing;
import brave.handler.SpanHandler;
import brave.sampler.Sampler;
import org.apache.dubbo.common.extension.ExtensionFactory;
import org.example.brave.rpc.RpcRequestMatchers;
import org.example.brave.rpc.RpcRuleSampler;
import org.example.brave.rpc.RpcTracing;
import zipkin2.reporter.AsyncReporter;
import zipkin2.reporter.brave.ZipkinSpanHandler;
import zipkin2.reporter.okhttp3.OkHttpSender;

public class TracingExtensionFactory implements ExtensionFactory {
    @Override
    public <T> T getExtension(Class<T> type, String name) {
        if (type != RpcTracing.class) return null;
        return (T) RpcTracing.newBuilder(tracing())
                .serverSampler(serverSampler())
                .build();
    }

    RpcRuleSampler serverSampler() {
        return RpcRuleSampler.newBuilder()
                .putRule(RpcRequestMatchers.methodEquals("sayHello"), Sampler.NEVER_SAMPLE)
                .build();
    }

    Tracing tracing() {
        return Tracing.newBuilder()
                .localServiceName("provider")
                .addSpanHandler(spanHandler())
                .build();
    }

    SpanHandler spanHandler() {
        return ZipkinSpanHandler.create(AsyncReporter.builder(sender()).build());
    }

    OkHttpSender sender() {
        return OkHttpSender.newBuilder().endpoint("http://localhost:9411/api/v2/spans").build();
    }
}
